import os
import logging
import numpy as np
import trimesh
import pyqtgraph.opengl as gl
from enum import Enum

logger = logging.getLogger('ModelManager')

class ModelType(Enum):
    """
    定义所有可用的模型类型
    """
    # --- 程序化生成的模型
    UAV_JET       = 0      # 简易三角翼战斗机模型
    UAV_PROPELLER = 1      # 简易螺旋桨飞机模型
    UAV_STEALTH   = 2      # 简易隐身无人机模型
    UAV_B2        = 3      # 简易B2风格无人机模型
    UAV_MISSILE   = 4      # 简易巡飞弹无人机模型

    # --- 外部模型文件
    UAV_CRAFT     = 100    # 宇宙飞船


class MeshBuilder:
    """
    一个简单的 3D 几何体构建辅助类。
    用于将简单的几何形状（方块、锥体）拼接成复杂的网格。
    """
    def __init__(self):
        self.verts = []
        self.faces = []
    
    def add_box(self, center, size):
        """添加一个长方体"""
        x, y, z = center
        dx, dy, dz = size[0]/2, size[1]/2, size[2]/2
        
        # 记录当前顶点起始索引
        start_idx = len(self.verts)
        
        # 8个顶点
        v = [
            [x-dx, y-dy, z-dz], [x+dx, y-dy, z-dz], [x+dx, y+dy, z-dz], [x-dx, y+dy, z-dz], # 底面 0-3
            [x-dx, y-dy, z+dz], [x+dx, y-dy, z+dz], [x+dx, y+dy, z+dz], [x-dx, y+dy, z+dz]  # 顶面 4-7
        ]
        self.verts.extend(v)
        
        # 12个面 (每个面由2个三角形组成)
        # 顺时针或逆时针并不太影响 pyqtgraph，因为它通常双面渲染
        f_indices = [
            [0,1,2], [0,2,3], # Bottom
            [4,5,6], [4,6,7], # Top
            [0,1,5], [0,5,4], # Front (Y-)
            [2,3,7], [2,7,6], # Back (Y+)
            [1,2,6], [1,6,5], # Right (X+)
            [3,0,4], [3,4,7]  # Left (X-)
        ]
        
        # 将局部索引转换为全局索引
        self.faces.extend([[idx + start_idx for idx in tri] for tri in f_indices])

    def add_pyramid(self, center, size, axis='x'):
        """添加一个简单的锥体 (用于机头或导弹头)"""
        x, y, z = center
        dx, dy, dz = size[0]/2, size[1]/2, size[2]/2
        start_idx = len(self.verts)
        
        # 简单起见，底面是矩形，顶点汇聚于一点
        # 这里假设 axis='x'，即尖端指向 X 正方向
    
        tip = [x+dx, y, z] # 尖端
        # 底面 4 点
        base = [
            [x-dx, y-dy, z-dz], 
            [x-dx, y+dy, z-dz], 
            [x-dx, y+dy, z+dz], 
            [x-dx, y-dy, z+dz]
        ]
        
        self.verts.append(tip)
        self.verts.extend(base)
        
        # 0是尖端，1-4是底面
        f_indices = [
            [1,2,3], [1,3,4], # 底面封闭
            [0,1,2], # 侧面1
            [0,2,3], # 侧面2
            [0,3,4], # 侧面3
            [0,4,1]  # 侧面4
        ]
        self.faces.extend([[idx + start_idx for idx in tri] for tri in f_indices])

    def get_mesh_data(self):
        """返回 pyqtgraph.opengl.MeshData 对象"""
        return gl.MeshData(vertexes=np.array(self.verts), faces=np.array(self.faces))

class ModelManager:
    """
    模型管理单例类
    负责生成或加载3D模型数据，并缓存以供复用
    """
    
    _instance = None
    _mesh_cache = {}  # 缓存字典: {ModelType: gl.MeshData}

    # 定义外部模型文件的路径映射
    # 请根据实际情况修改文件路径
    _model_config = {
        ModelType.UAV_CRAFT : {
            'path': 'resource/3DModels/craft_speederC.obj',
            # 绕Z轴旋转90度
            'rotate_angle': 90,
            'rotate_axis': [1, 0, 0],
            'target_faces': 300,
            'scale': 1.0
        }
    }

    @classmethod
    def get_mesh_data(cls, model_type: ModelType, color=None):
        """
        获取模型数据的核心方法
        :param model_type: 模型类型
        :param color: (可选) 统一覆盖颜色 (r,g,b,a)
        :return: gl.MeshData 对象
        """
        # 1. 检查缓存(性能优化)
        if model_type in cls._mesh_cache:
            # 如果需要修改颜色，这里可能需要深拷贝
            # 但为了性能通常建议在创建时设置color，而不是修改MeshData
            return cls._mesh_cache[model_type]
    
        logger.info(f'正在初始化模型数据: {model_type.name} ...')

        md = None

        # 2. 根据类型分发处理
        if model_type == ModelType.UAV_JET:
            md = cls._create_uav_jet()
        elif model_type == ModelType.UAV_PROPELLER:
            md = cls._create_uav_propeller()
        elif model_type == ModelType.UAV_STEALTH:
            md = cls._create_uav_stealth()
        elif model_type == ModelType.UAV_B2:
            md = cls._create_uav_b2()
        elif model_type == ModelType.UAV_MISSILE:
            md = cls._create_uav_missile()
        elif model_type == ModelType.UAV_CRAFT:
            cfg = cls._model_config[model_type]
            print(cfg)
            md = cls._load_obj_file(
                cfg['path'], 
                cfg['rotate_angle'], 
                cfg['rotate_axis'],
                cfg['target_faces'], 
                cfg['scale'])
        else:
            logger.error(f'错误：未知的模型类型 {model_type}')
            return None
        
        # 3. 后处理：计算法线 (解决光照下模型发黑的问题)
        # 这一步非常重要，否则 shader='shaded' 会失效
        if md is not None:
            # 如果没有顶点法线，计算一下
            # if md.vertexNormals() is None:
            #     md.computeVertexNormals()
            
            # 存入缓存
            cls._mesh_cache[model_type] = md
        
        return md
    
    @staticmethod
    def _create_uav_jet():
        """生成的战斗机模型 (三角翼)"""
        mb = MeshBuilder()
        # 机身 (长条)
        mb.add_box((0, 0, 0), (1.5, 0.25, 0.25))
        # 进气道/机身加宽
        mb.add_box((-0.2, 0, 0), (0.8, 0.5, 0.2))
        # 机翼 (向后掠的扁平块)
        # 用多个 Box 模拟三角翼形状
        mb.add_box((-0.3, 0, 0), (0.6, 1.4, 0.05)) # 主翼
        # 垂直尾翼
        mb.add_box((-0.6, 0, 0.3), (0.4, 0.05, 0.6))
        # 平尾
        mb.add_box((-0.7, 0, 0), (0.3, 0.8, 0.05))
        return mb.get_mesh_data()
    
    @staticmethod
    def _create_uav_propeller():
        """生成的螺旋桨飞机 (类似塞斯纳)"""
        mb = MeshBuilder()
        # 机身
        mb.add_box((0, 0, 0), (1.2, 0.3, 0.3))
        # 机尾变细的部分 (简化为连接的一个较小的Box)
        mb.add_box((-0.8, 0, 0.05), (0.6, 0.15, 0.15))
        # 主翼 (上单翼)
        mb.add_box((0.2, 0, 0.2), (0.3, 1.8, 0.05))
        # 垂直尾翼
        mb.add_box((-1.0, 0, 0.2), (0.3, 0.05, 0.4))
        # 水平尾翼
        mb.add_box((-1.0, 0, 0.05), (0.25, 0.8, 0.05))
        # 螺旋桨 (十字)
        mb.add_box((0.6, 0, 0), (0.02, 0.8, 0.05))
        mb.add_box((0.6, 0, 0), (0.02, 0.05, 0.8))
        return mb.get_mesh_data()
    
    @staticmethod
    def _create_uav_stealth():
        """生成的飞翼/隐身无人机"""
        mb = MeshBuilder()
        # 中心机身
        mb.add_box((0, 0, 0), (0.8, 0.4, 0.15))
        # 机翼 (大后掠角) - 这里用阶梯状的 Box 模拟 V 形
        mb.add_box((-0.2, 0.6, 0), (0.6, 0.8, 0.08)) # 左内
        mb.add_box((-0.2, -0.6, 0), (0.6, 0.8, 0.08)) # 右内
        mb.add_box((-0.5, 1.2, 0), (0.4, 0.6, 0.05)) # 左外
        mb.add_box((-0.5, -1.2, 0), (0.4, 0.6, 0.05)) # 右外
        return mb.get_mesh_data()
    
    @staticmethod
    def _create_uav_b2():
        """
        生成一个改进版的 B2 隐身轰炸机模型。
        特点：流线型机身、W型锯齿尾部、融合翼身。
        """
        # --- 1. 定义关键顶点坐标 (x, y, z) ---
        # 坐标系: X朝前, Y朝左, Z朝上
        
        # 中轴线上的点
        nose = [1.5, 0.0, 0.0]        # 0: 机头
        cockpit_top = [0.2, 0.0, 0.25]# 1: 驾驶舱顶部隆起
        tail_center = [-0.8, 0.0, 0.05] # 2: 机尾中心
        belly_center = [0.0, 0.0, -0.1] # 3: 机腹中心 (略微下凸)
        
        # 左翼关键点 (Y > 0)
        # 前缘转折点 (Leading Edge Break)
        le_mid = [0.0, 1.2, 0.05]     # 4
        # 翼尖 (Wing Tip)
        wing_tip = [-1.5, 3.5, 0.1]   # 5
        # 后缘锯齿点 (Trailing Edge Zig-Zags) - 制造W型尾
        te_out = [-1.8, 2.0, 0.05]    # 6: 外侧锯齿凹陷
        te_in =  [-1.2, 0.8, 0.05]    # 7: 内侧锯齿凸起

        # --- 2. 镜像生成右翼点 (Y < 0) ---
        # 对应索引 8, 9, 10, 11
        right_points = []
        for p in [le_mid, wing_tip, te_out, te_in]:
            right_points.append([p[0], -p[1], p[2]])
            
        # 合并所有顶点
        # 索引映射:
        # 0:Nose, 1:Top, 2:Tail, 3:Belly
        # 4:L_LE, 5:L_Tip, 6:L_TE_Out, 7:L_TE_In
        # 8:R_LE, 9:R_Tip, 10:R_TE_Out, 11:R_TE_In
        verts = np.array([nose, cockpit_top, tail_center, belly_center, 
                          le_mid, wing_tip, te_out, te_in] + right_points)

        # --- 3. 定义面 (Faces) ---
        # 使用三角形连接顶点。注意逆时针顺序(CCW)通常为正面，但PyQtGraph双面渲染，主要影响光照法线。
        faces = [
            # --- 上表面 (Top Surface) ---
            # 左侧机身与机翼
            [0, 1, 4],    # 机头-座舱-前缘
            [1, 4, 7],    # 座舱-前缘-后缘内
            [1, 7, 2],    # 座舱-后缘内-机尾
            [4, 5, 6],    # 前缘-翼尖-后缘外
            [4, 6, 7],    # 前缘-后缘外-后缘内
            
            # 右侧机身与机翼 (镜像)
            [0, 8, 1],    # 机头-前缘-座舱
            [1, 11, 8],   # 座舱-后缘内-前缘
            [1, 2, 11],   # 座舱-机尾-后缘内
            [8, 10, 9],   # 前缘-后缘外-翼尖
            [8, 11, 10],  # 前缘-后缘内-后缘外

            # --- 下表面 (Bottom Surface / Belly) ---
            # 为了简化，下表面做得比较平坦
            # 左侧
            [0, 4, 3],    # 机头-前缘-机腹
            [3, 4, 7],    # 机腹-前缘-后缘内
            [3, 7, 2],    # 机腹-后缘内-机尾
            [4, 5, 6],    # 前缘-翼尖-后缘外 (下)
            [4, 6, 7],    # 前缘-后缘外-后缘内 (下)
            
            # 右侧
            [0, 3, 8],
            [3, 11, 8],
            [3, 2, 11],
            [8, 9, 10],
            [8, 10, 11]
        ]

        faces = np.array(faces)
        
        # --- 4. 生成 MeshData 并计算法线 ---
        md = gl.MeshData(vertexes=verts, faces=faces)
        return md
    
    @staticmethod
    def _create_uav_missile():
        """生成的巡飞弹"""
        mb = MeshBuilder()
        # 弹体
        mb.add_box((0, 0, 0), (2.0, 0.2, 0.2))
        # 弹头 (锥体) - 这里用pyqtgraph自带的 Cylinder 可能更圆，但为了统一这里用 Box 模拟或自定义 Pyramid
        mb.add_pyramid((1.1, 0, 0), (0.4, 0.2, 0.2), axis='x')
        # 尾翼 (X型)
        fin_size = 0.6
        mb.add_box((-0.8, 0, 0), (0.3, fin_size, 0.05)) # 水平
        mb.add_box((-0.8, 0, 0), (0.3, 0.05, fin_size)) # 垂直
        return mb.get_mesh_data()
    
    @staticmethod
    def _load_obj_file(file_path, rotate_angle=0, rotate_axis=None, target_faces=300, scale=1.0):
        """
        加载 STL/OBJ 模型，减少面数，并转换为 pyqtgraph 可用的 MeshData
        :param file_path: 模型路径 (.obj 或 .stl)
        :param rotate_angle: 旋转角度 (单位：度)
        :param rotate_axis: 旋转轴
        :param target_faces: 目标面数 (越低性能越好，推荐 100-500)
        :param scale: 缩放比例
        :return: gl.MeshData
        """
        # 1. 使用 trimesh 加载模型
        mesh = trimesh.load(file_path)
        
        # 如果是包含多个物体的场景，合并成一个 mesh
        if isinstance(mesh, trimesh.Scene):
            mesh = trimesh.util.concatenate(mesh.dump())

        # 2. 应用旋转
        if rotate_angle != 0 and rotate_axis is not None:
            from trimesh import transformations
            # 将角度转换为弧度
            rad = np.radians(rotate_angle)
            # 生成旋转矩阵 (4x4)
            matrix = transformations.rotation_matrix(rad, rotate_axis)
            # 应用变换
            mesh.apply_transform(matrix)

        # 3. 减面 (Decimation) - 这一步对性能至关重要
        # 如果当前面数已经很少，就不减了
        if len(mesh.faces) > target_faces:
            print(f"原始面数: {len(mesh.faces)}，正在简化为 {target_faces}...")
            # 这种简化算法会保持物体的大致形状
            mesh = mesh.simplify_quadratic_decimation(target_faces)
            print(f"简化后面数: {len(mesh.faces)}")
        
        # 4. 提取顶点和面
        # trimesh 的顶点通常是 (N, 3) 的 float 数组
        verts = mesh.vertices * scale
        faces = mesh.faces

        # 4. 生成 pyqtgraph 的 MeshData
        # 计算法线可以让光照效果更好
        mesh_data = gl.MeshData(vertexes=verts, faces=faces)
        
        return mesh_data